// do.go
package main

import (
	"crypto/rand"
	"log"
	"math/big"
	"strconv"

	pbc "github.com/Nik-U/pbc"
)

// DataOwner 代表数据所有者
type DataOwner struct {
	sk           *SecretKey
	pk           *PublicKey
	params       *SystemParams
	originalData map[string][]*pbc.Element // 存储原始数据块 (Z_p 元素)
	totalBlocks  int                       // 文件总块数
}

// NewDataOwner 创建一个新的数据所有者，并完成 KeyGen
func NewDataOwner(params *SystemParams) (*DataOwner, error) {
	// KeyGen
	gamma := params.Pairing.NewZr().Rand()
	g_gamma := params.Pairing.NewG1().PowZn(params.G, gamma)

	log.Println("[DO] Generated private key gamma and public key g^gamma.")

	return &DataOwner{
		sk:           &SecretKey{Gamma: gamma},
		pk:           &PublicKey{G_gamma: g_gamma},
		params:       params,
		originalData: make(map[string][]*pbc.Element),
	}, nil
}

func (do *DataOwner) TagGen(fileData []byte, fileID string, blockSize int, cspIDs []string) (map[string][]*pbc.Element, error) { // 接收 cspIDs
	log.Printf("[DO] Starting TagGen for file '%s'...", fileID)
	// 1. 分块 (不变)
	blocks := make([]*pbc.Element, 0)
	for i := 0; i < len(fileData); i += blockSize {
		end := i + blockSize
		if end > len(fileData) {
			end = len(fileData)
		}
		blockElement := do.params.Pairing.NewZr().SetFromHash(fileData[i:end])
		blocks = append(blocks, blockElement)
	}
	do.originalData[fileID] = blocks
	do.totalBlocks = len(blocks)
	log.Printf("[DO] File split into %d blocks.", do.totalBlocks)

	// 2. 为每个CSP的每个块生成标签 y_ki
	tagMap := make(map[string][]*pbc.Element)

	for _, cspID := range cspIDs { // 使用传入的 cspIDs
		tags := make([]*pbc.Element, do.totalBlocks)
		for i, xi := range blocks {
			h_ki := do.params.h(cspID + "||" + fileID + "||" + strconv.Itoa(i))
			u_xki := do.params.Pairing.NewG1().PowZn(do.params.U, xi)
			base := do.params.Pairing.NewG1().Mul(h_ki, u_xki)
			tags[i] = do.params.Pairing.NewG1().PowZn(base, do.sk.Gamma)
		}
		tagMap[cspID] = tags
	}
	log.Printf("[DO] TagGen finished. Generated tags for %d CSPs.", len(cspIDs))
	return tagMap, nil
}

// GenChallenge 核心算法：生成审计挑战
func (do *DataOwner) GenChallenge(fileID string, challengedBlocks int) (*AuditChallenge, error) {
	indices := make([]int, challengedBlocks)
	weights := make([]*pbc.Element, challengedBlocks)

	for i := 0; i < challengedBlocks; i++ {
		idx, _ := rand.Int(rand.Reader, big.NewInt(int64(do.totalBlocks)))
		indices[i] = int(idx.Int64())
		// 权重 e_l 是 Z_p 中的随机元素
		weights[i] = do.params.Pairing.NewZr().Rand()
	}

	return &AuditChallenge{
		FileID:  fileID,
		Indices: indices,
		Weights: weights,
	}, nil
}
